/*
* This file is a part of the open-eBackup project.
* This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
* If a copy of the MPL was not distributed with this file, You can obtain one at
* http://mozilla.org/MPL/2.0/.
*
* Copyright (c) [2024] Huawei Technologies Co.,Ltd.
*
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
* EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
* MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
*/

#define WIN32_LEAN_AND_MEAN
#include <objbase.h>
#include <process.h>
#include <iostream>
#include <string>
#include <cstdio>
#include <atlstr.h>
#include "json/json.h"
#include "vdi.h"
#include "log/Log.h"
#include "common/Path.h"
#include "common/Utils.h"
#include "config_reader/ConfigIniReader.h"

namespace {
    const int SQLSERVE_FULL_BACKUP = 1;
    const int SQLSERVE_DIFF_BACKUP = 2;
    const int SQLSERVE_LOG_BACKUP = 3;
    const int SQLSERVE_END_LOG_BACKUP = 4;
    const int SQLSERVE_FULL_RESTORE_WITH_REPLACE = 5;
    const int SQLSERVE_FULL_RESTORE_WITH_NORECOVERY = 6;
    const int SQLSERVE_DIFF_RESTORE_WITH_RECOVERY = 7;
    const int SQLSERVE_DIFF_RESTORE_WITH_NORECOVERY = 8;
    const int SQLSERVE_LOG_RESTORE_WITH_NORECOVERY = 9;
    const int SQLSERVE_LOG_RESTORE_WITH_RECOVERY = 10;

    const int SQLSERVE_CONFIGURATION_NUM = 120000;
    const int SQLSERVE_VDS_WNAME_SIZE = 49;
    const int SQLSERVE_VDS_NAME_SIZE = 100;
    const int CMD_LENG = 1024;

    const int CONSTANTS_0 = 0;
    const int CONSTANTS_1 = 1;
    const int CONSTANTS_2 = 2;
    const int CONSTANTS_3 = 3;
    const int CONSTANTS_4 = 4;
    const int CONSTANTS_5 = 5;
    const int CONSTANTS_6 = 6;
    const int CONSTANTS_7 = 7;
    const int CONSTANTS_8 = 8;
    const int CONSTANTS_256 = 256;
    int g_threadNum = 4;         // g_threadNum 默认为4
    const int BACKUP_TIME_OUT = 2 * 60 * 60 * 1000; // 数据库备份恢复超时时间2小时 (测试速度为200M/S，2小时能备份恢复700G)
    const int VDF_WRITE_MEDIA  = 0x10000;

    WCHAR wVdsName[50];
    IClientVirtualDeviceSet2* g_vds = nullptr;
    VDConfig g_config;

struct SQLserverInputInfo {
    std::string dbName; // 数据库名称
    std::string cifsPath; // 备份生成的备份文件路径（即挂载的cifs的文件路径）
    std::string userInfo; // 健全信息
    std::string backupTime; // 按时间点恢复的时间点信息
    std::string restoreInfo; // 回复到新位置的目标位置信息
    int actionType; // 操作类型 1：全量备份 2：差异备份 3：日志备份 4 尾日志备份 5 全量恢复 6 日志的差异、全量恢复 7 差异的全量恢复 8 差异恢复  9按时间点恢复  10 按时间点的尾恢复
    int parallelProcess;    // 并发数 不传默认为4， 就是备份的时候起多少个线程
};

struct THREAD_INFO {
    IClientVirtualDeviceSet *vds;
    int threadId;
    std::string fileName; // 数据库名称
};

bool PerformTransfer(IClientVirtualDevice* vd, int backup, const std::string &backupPathName)
{
    VDC_Command*   cmd;
    DWORD           completionCode;
    DWORD           bytesTransferred;
    HRESULT         hr;
    int             termCode = -1;

    FILE* fh = fopen(backupPathName.c_str(), (backup) ? "wb" : "rb");
    if (fh == nullptr) {
        printf("Failed to open: %s\n", backupPathName.c_str());
        return false;
    }

    while (SUCCEEDED(hr = vd->GetCommand(INFINITE, &cmd))) {
        bytesTransferred = 0;
        switch (cmd->commandCode) {
            case VDC_READ:
                bytesTransferred = fread(cmd->buffer, 1, cmd->size, fh);
                break;

            case VDC_WRITE:
                bytesTransferred = fwrite(cmd->buffer, 1, cmd->size, fh);
                break;

            case VDC_FLUSH:
                fflush(fh);
                completionCode = ERROR_SUCCESS;
                break;

            case VDC_CLEAR_ERROR:
                completionCode = ERROR_SUCCESS;
                break;

            default:
                completionCode = ERROR_NOT_SUPPORTED;
                break;
        }
        if (bytesTransferred == cmd->size) {
            completionCode = ERROR_SUCCESS;
        } else {
            completionCode = ERROR_DISK_FULL;
        }

        hr = vd->CompleteCommand(cmd, completionCode, bytesTransferred, 0);
        if (!SUCCEEDED(hr)) {
            printf("Completion Failed: x%X\n", hr);
            break;
        }
    }

    if (hr != VD_E_CLOSE) {
        printf("Unexpected termination: x%X\n", hr);
    } else {
        termCode = 0;
    }
    fclose(fh);
    return termCode;
}

std::string GetVDSName()
{
    std::string vdsName;
    int vdsNamesize = WideCharToMultiByte(CP_ACP, 0, wVdsName, wcslen(wVdsName), nullptr, 0, nullptr, nullptr);
    if (vdsNamesize < 0) {
        return "";
    }
    std::shared_ptr<char[]> cVdsName = std::make_unique<char[]>(vdsNamesize + 1);
    WideCharToMultiByte(CP_ACP, 0, wVdsName, wcslen(wVdsName) + 1, cVdsName.get(), vdsNamesize, nullptr, nullptr);
    vdsName = cVdsName.get();

    std::string fullVDSInfo;
    for (int i = 0; i < g_threadNum; i++) {
        if (i == 0) {
            fullVDSInfo += " VIRTUAL_DEVICE='" + vdsName + "' ";
        } else {
            fullVDSInfo += ", VIRTUAL_DEVICE='" + vdsName + std::to_string(i) + "' ";
        }
    }
    printf("fullVDSInfo: %s\n", fullVDSInfo.c_str());
    return fullVDSInfo;
}

std::string Getsqlcmd(const SQLserverInputInfo &taskInfo, const std::string &vdsName)
{
    std::string sqlcmd;
    switch (taskInfo.actionType) {
        case SQLSERVE_FULL_BACKUP:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"BACKUP DATABASE [" + taskInfo.dbName + "] TO " + vdsName;
            break;
        case SQLSERVE_DIFF_BACKUP:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"BACKUP DATABASE [" + taskInfo.dbName + "] TO " +
                vdsName + " WITH DIFFERENTIAL";
            break;
        case SQLSERVE_LOG_BACKUP:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"BACKUP LOG [" + taskInfo.dbName + "] TO " + vdsName;
            break;
        case SQLSERVE_END_LOG_BACKUP:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"BACKUP LOG [" + taskInfo.dbName + "] TO " +
                vdsName + " WITH NORECOVERY";
            break;
        case SQLSERVE_FULL_RESTORE_WITH_REPLACE:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"RESTORE DATABASE [" + taskInfo.dbName + "] FROM " +
                vdsName + " WITH REPLACE";
            break;
        case SQLSERVE_FULL_RESTORE_WITH_NORECOVERY:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"RESTORE DATABASE [" + taskInfo.dbName + "] FROM " +
                vdsName + " WITH NORECOVERY";
            break;
        case SQLSERVE_DIFF_RESTORE_WITH_RECOVERY:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"RESTORE DATABASE [" + taskInfo.dbName + "] FROM " +
                vdsName + " WITH REPLACE,NORECOVERY";
            break;
        case SQLSERVE_DIFF_RESTORE_WITH_NORECOVERY:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"RESTORE DATABASE [" + taskInfo.dbName + "] FROM " +
                vdsName + " WITH RECOVERY";
            break;
        case SQLSERVE_LOG_RESTORE_WITH_NORECOVERY:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"RESTORE LOG [" + taskInfo.dbName + "] FROM " +
                vdsName + " WITH NORECOVERY,STOPAT='" + taskInfo.backupTime + "'";
            break;
        case SQLSERVE_LOG_RESTORE_WITH_RECOVERY:
            sqlcmd = "sqlcmd " + taskInfo.userInfo + " -b -Q\"RESTORE DATABASE [" + taskInfo.dbName +
                "] WITH RECOVERY ";
            break;
        default:
            printf("error input data: %d\n", taskInfo.actionType);
            return "";
    }
    sqlcmd =  sqlcmd + taskInfo.restoreInfo + "\"";
    INFOLOG("sqlCmd is : %s", sqlcmd.c_str());
    return sqlcmd;
}

HANDLE ExecSQL(const SQLserverInputInfo &taskInfo)
{
    PROCESS_INFORMATION pi;
    STARTUPINFO si;

    std::string vdsName = GetVDSName();
    std::string sqlcmd = Getsqlcmd(taskInfo, vdsName);
    if (sqlcmd.empty()) {
        return nullptr;
    }
    int cmdLength = strlen(sqlcmd.c_str());
    std::shared_ptr<char[]> wVdsInfo = std::make_unique<char[]>(cmdLength + 1);
    if (memcpy_s(wVdsInfo.get(), strlen(sqlcmd.c_str()) + 1, sqlcmd.c_str(), strlen(sqlcmd.c_str()) + 1) != 0) {
        printf("MemCopy of vdsInfo failed.\n");
        return nullptr;
    }

    GetStartupInfo(&si);
    bool ret = CreateProcess(NULL, wVdsInfo.get(), NULL, NULL, TRUE, 0,  NULL, NULL, &si, &pi);
    if (!ret) {
        printf("CreateProcess failed, errcode:[%d]\n", GetLastError());
        return nullptr;
    }

    return pi.hProcess;
}

void StrSplit(std::vector<std::string>& vecTokens, const std::string& strText, char cSep)
{
    int start = 0;
    int end;
    while ((end = strText.find(cSep, start)) != std::string::npos) {
        vecTokens.push_back(strText.substr(start, end - start));
        start = end + 1;
    }

    vecTokens.push_back(strText.substr(start));
}

bool GetVDSInfo(const std::string& userInfo)
{
    HRESULT hr = CoCreateInstance(IID_IClientVirtualDeviceSet, nullptr, CLSCTX_INPROC_SERVER,
        IID_IClientVirtualDeviceSet2, (void**)&g_vds);
    if (!SUCCEEDED(hr)) {
        printf("Could not create component: x%X, Check registration of SQLVDI.DLL and value of IID \n", hr);
        return false;
    }

    GUID vdsId;
    CoCreateGuid(&vdsId);
    StringFromGUID2(vdsId, wVdsName, SQLSERVE_VDS_WNAME_SIZE);
    g_config.deviceCount = g_threadNum;

    std::string instanceName = "";
    std::vector<std::string> vecUserInfo;
    StrSplit(vecUserInfo, userInfo, ' ');
    if (vecUserInfo.size() >= CONSTANTS_2) {
        std::vector<std::string> vecinstanceName;
        StrSplit(vecinstanceName, vecUserInfo[CONSTANTS_1], '\\');
        if (vecinstanceName.size() >= CONSTANTS_2) {
            instanceName = vecinstanceName[CONSTANTS_1];
        }
    }

    if (instanceName.empty()) {
        hr = g_vds->CreateEx(nullptr, wVdsName, &g_config);
    } else { // 对于没有默认实例的场景，需要根据实例名来创建虚拟磁盘
        WCHAR lpInstanceName[50];
        int cmdLength = strlen(instanceName.c_str());
        MultiByteToWideChar(CP_ACP, 0, instanceName.c_str(), cmdLength + 1, lpInstanceName, cmdLength + 1);
        hr = g_vds->CreateEx(lpInstanceName, wVdsName, &g_config);
    }

    if (!SUCCEEDED(hr)) {
        printf("VDS::Create fails: x%X", hr);
        return false;
    }
    return true;
}

unsigned __stdcall RunSecondary(LPVOID threadInfo)
{
    HRESULT hr;
    WCHAR vdsName[SQLSERVE_VDS_NAME_SIZE];
    IClientVirtualDevice* vd;
    VDConfig threadConfig;
    int termCode = -1;

    int threadId = ((THREAD_INFO*)threadInfo)->threadId;
    IClientVirtualDeviceSet *vdsHandle = ((THREAD_INFO*)threadInfo)->vds;
    std::string backupPathName = ((THREAD_INFO*)threadInfo)->fileName;

    if (threadId == 0) {
        int ret = wcscpy_s(vdsName, SQLSERVE_VDS_NAME_SIZE, wVdsName);
        if (ret != 0) {
            printf("Used wcscpy_s failed.\n");
            return termCode;
        }
    } else {
        int ret = swprintf_s(vdsName, SQLSERVE_VDS_NAME_SIZE, L"%s%d", wVdsName, threadId);
        if (ret == -1) {
            printf("Used swprintf_s failed.\n");
            return termCode;
        }
    }

    hr = vdsHandle->OpenDevice(vdsName, &vd);
    if (!SUCCEEDED (hr)) {
        printf("OpenDevice fails on %ls: x%X", vdsName, hr);
        vdsHandle->SignalAbort();
        return termCode;
    }

    hr = vdsHandle->GetConfiguration(INFINITE, &threadConfig);
    if (!SUCCEEDED (hr)) {
        printf("thread Getconfig fails: x%X\n", hr);
        vdsHandle->SignalAbort();
        return termCode;
    }

    printf("Performing data transfer...\n");
    termCode = PerformTransfer(vd, (threadConfig.features & VDF_WRITE_MEDIA), backupPathName);
    vd->Release();
    if (termCode != 0) {
        printf("Performing data transfer failed.\n");
        vdsHandle->SignalAbort();
    }
    return termCode;
}

bool AddThreadInfo(std::vector<THREAD_INFO>& threadInfo, SQLserverInputInfo taskInfo,
    std::vector<HANDLE>& childrenInfo)
{
    unsigned threadId;
 
    for (int i = 0; i < g_threadNum; i++) {
        threadInfo[i].vds = g_vds;
        threadInfo[i].threadId = i;
        if (i == 0) {
            threadInfo[i].fileName = taskInfo.cifsPath;
        } else {
            threadInfo[i].fileName = taskInfo.cifsPath + std::to_string(i);
        }
        childrenInfo[i] = (HANDLE)_beginthreadex(nullptr, 0, RunSecondary, (void*)&threadInfo[i], 0, &threadId);
 
        if (childrenInfo[i] == nullptr) {
            printf("Failed to create thread. errinfo is [%d]\n", errno);
            g_vds->SignalAbort();
            return false;
        }
    }
    return true;
}

void WaitSubThreadComplete(std::vector<HANDLE> &childrenInfo, bool &threadExeResult)
{
    DWORD waitStatus;
    DWORD exitCode;
    waitStatus = WaitForMultipleObjects(g_threadNum, childrenInfo.data(), TRUE, INFINITE);
    if (waitStatus == WAIT_FAILED) {
        printf("WaitForMultipleObjects failed. Error: %s\n", GetLastError());
        threadExeResult = false;
    } else if (waitStatus < WAIT_OBJECT_0 || waitStatus >= WAIT_OBJECT_0 + g_threadNum) {
        printf("WaitForMultipleObjects failed, errcode:[%d] errorInfo:[%d].\n", GetLastError(), waitStatus);
        threadExeResult = false;
    }
    for (int i = 0; i < g_threadNum; ++i) {
        GetExitCodeThread(childrenInfo[i], &exitCode);
        printf("Thread exit code:[%d]\n", exitCode);
        if (exitCode != 0) {
            threadExeResult = false;
        }
        CloseHandle(childrenInfo[i]);
    }
}

void WaitSubProcessComplete(HANDLE &hSQLProcess, bool &processExeResult)
{
    DWORD waitStatus;
    DWORD exitCode;
    waitStatus = WaitForSingleObject(hSQLProcess, INFINITE);
    if (waitStatus == WAIT_FAILED) {
        printf("WaitForSingleObject failed. Error: %s\n", GetLastError());
        processExeResult = false;
    } else if (waitStatus != WAIT_OBJECT_0) {
        printf("WaitForSingleObject failed, errcode:[%d] errorInfo:[%d].\n", GetLastError(), waitStatus);
        processExeResult = false;
    }
    GetExitCodeProcess(hSQLProcess, &exitCode);
    printf("Process exit code:[%d]\n", exitCode);
    if (exitCode != 0) {
        processExeResult = false;
    }
    CloseHandle(hSQLProcess);
}

bool WriteInfoToVDS(SQLserverInputInfo taskInfo, HANDLE hSQLProcess)
{
    bool threadExeResult = true;
    bool processExeResult = true;
    HRESULT hr = g_vds->GetConfiguration(SQLSERVE_CONFIGURATION_NUM, &g_config);
    if (!SUCCEEDED(hr)) {
        printf("GetConfiguration fails: errInfo:x%X\n", hr);
        return false;
    }

    std::vector<THREAD_INFO> threadInfo {};
    std::vector<HANDLE> childrenInfo {};
    for (int i = 0; i < g_threadNum; i++) {
        THREAD_INFO thread_info;
        threadInfo.emplace_back(thread_info);
    }

    for (int i = 0; i < g_threadNum; i++) {
        HANDLE _handle;
        childrenInfo.emplace_back(_handle);
    }

    if (!AddThreadInfo(threadInfo, taskInfo, childrenInfo)) {
        return false;
    }

    printf("All children thread running. Waiting for completion.\n");
    WaitSubThreadComplete(childrenInfo, threadExeResult);
    // 等待sqlcmd进程结束
    WaitSubProcessComplete(hSQLProcess, processExeResult);

    g_vds->SignalAbort();
    if (threadExeResult && processExeResult) {
        printf("Thread and process exec succ.\n");
        return true;
    }
    printf("Thread and process exec failed.\n");
    return false;
}

bool JsonStringToJsonValue(const std::string& jsonString, Json::Value& value)
{
    if (jsonString.empty()) {
        printf("jsonString is empty.");
        return false;
    }
    Json::CharReaderBuilder charReaderBuilder;
    std::unique_ptr<Json::CharReader> pCharReader(charReaderBuilder.newCharReader());
    if (!pCharReader) {
        printf("CharReaderBuilder is null.");
        return false;
    }
    const char* begin = jsonString.c_str();
    const char* end = begin + jsonString.size();
    std::string strError;
    bool ret = false;
    try {
        ret = pCharReader->parse(begin, end, &value, &strError);
    }
    catch (std::exception& e) {
        printf("failed to parse, throw a exception. %s", e.what());
        return false;
    }
    if (!ret) {
        printf("Json parse failed, error: %s", strError.c_str());
        return false;
    }
    return true;
}

void GetrestoreInfo(std::string &jsonString, std::string& value)
{
    Json::Value jsonData;
    JsonStringToJsonValue(jsonString, jsonData);

    for (int i = 0; i < jsonData.size(); i++) {
        value += ",move '" + jsonData[i]["name"].asString() + "' to '"
            + jsonData[i]["physical_name"].asString() + "' ";
    }
}

void InitLog()
{
    CString strPath;
    // C:\\DataBackup\\ProtectClient\\Plugins\\GeneralDBPlugin\\bin\\backupAPI.exe
 
    GetModuleFileName(NULL, strPath.GetBufferSetLength(MAX_PATH+1), MAX_PATH+1);
 
    std::string binPath((LPCTSTR)strPath);
    int rootPos = binPath.rfind("bin");
    std::string rootPath = binPath.substr(0, rootPos);
    Module::CPath::GetInstance().SetRootPath(rootPath);
 
    int logPos = rootPath.rfind("Plugins");
    std::string logFilePath = rootPath.substr(0, logPos) + "ProtectClient-E\\log\\Plugins\\GeneralDBPlugin";
    int logLevel = Module::ConfigReader::getInt("General", "LogLevel");
    int logCount = Module::ConfigReader::getInt("General", "LogCount");
    int logMaxSize = Module::ConfigReader::getInt("General", "LogMaxSize");
    Module::ConfigReader::getInt("General", "LogLevel");
    std::string logFile = "backup_api.log";
 
    Module::CLogger::GetInstance().Init(logFile.c_str(), logFilePath);
    Module::CLogger::GetInstance().SetLogConf(logLevel, logCount, logMaxSize);
}
}

int main(int argc, char *argv[])
{
    HANDLE hProcess;
    int ret = -1;

    if (argc < CONSTANTS_5) {
        printf("Input info invalid.\n");
        return ret;
    }
    InitLog();
    SQLserverInputInfo taskInfo;
    taskInfo.actionType = Module::SafeStoi(argv[1]);
    taskInfo.dbName = argv[CONSTANTS_2];
    taskInfo.cifsPath = argv[CONSTANTS_3];
    taskInfo.userInfo = argv[CONSTANTS_4];
    
    if (argc >= CONSTANTS_6) {
        taskInfo.backupTime = argv[CONSTANTS_5];
    }
    if (argc >= CONSTANTS_7) {
        taskInfo.restoreInfo = argv[CONSTANTS_6];
    }
    if (argc >= CONSTANTS_8) {
        taskInfo.parallelProcess = Module::SafeStoi(argv[CONSTANTS_7]);
        if (taskInfo.parallelProcess > CONSTANTS_0 && taskInfo.parallelProcess <= CONSTANTS_256) {
            g_threadNum = taskInfo.parallelProcess;
        }
    }
    
    HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    if (!SUCCEEDED(hr)) {
        printf("Coinit fails: x%X\n", hr);
        return ret;
    }
    if (GetVDSInfo(taskInfo.userInfo)) {
        hProcess = ExecSQL(taskInfo);
        if (hProcess == nullptr) {
            printf("ExecSQL failed.\n");
        } else if (taskInfo.actionType == SQLSERVE_LOG_RESTORE_WITH_RECOVERY) { // 按时间点恢复最后一次不需要写入读写数据
            printf("ExecSQL success.\n");
            ret = 0;
        } else {
            if (WriteInfoToVDS(taskInfo, hProcess)) {
                ret = 0;
            } else {
                printf("WriteInfoToVDS failed.\n");
            }
        }
        g_vds->Close();
        g_vds->Release();
    }

    CoUninitialize();
    return ret;
}

